package com.shujia.sql


import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

object Test5 {
  def main(args: Array[String]): Unit = {
    /**
     * 创建SparkSession的环境对象
     */
    val sparkSession: SparkSession = SparkSession.builder()
      .master("local")
      .appName("蚂蚁森林案例")
      .config("spark.sql.shuffle.partitions", "1")
      .getOrCreate()

    import org.apache.spark.sql.functions._
    import sparkSession.implicits._

    //读取用户每日碳排放量信息表
    val userLowCarbonDF: DataFrame = sparkSession.read
      .format("csv")
      .option("sep", "\t")
      .schema("user_id STRING,date_dt STRING,low_carbon Double")
      .load("spark/data/ant_user_low_carbon.txt")

    val plantCarbonDF: DataFrame = sparkSession.read
      .format("csv")
      .option("sep", "\t")
      .schema("plant_id STRING,plant_name STRING,plant_carbon Double")
      .load("spark/data/ant_plant_carbon.txt")

    //因为用户信息表与植物信息表是没有直接关联条件的，需要单独的从植物信息表中将胡杨和沙柳的所需能量提取出来由变量保存
    val huYangCarbon: Double = plantCarbonDF.where($"plant_name" === "胡杨")
      .select($"plant_carbon")
      .rdd
      .collect()
      .head
      .getAs[Double]("plant_carbon")

    val shaLiuCarbon: Double = plantCarbonDF.where($"plant_name" === "沙柳")
      .select($"plant_carbon")
      .rdd
      .collect()
      .head
      .getAs[Double]("plant_carbon")

    println(s"胡杨所需碳排放量：$huYangCarbon, 沙柳所需碳排放量：$shaLiuCarbon")
    println("==========================================================================")

    /**
     * 题目一：蚂蚁森林植物申领统计
     *    假设2017年1月1日开始记录低碳数据（user_low_carbon），假设2017年10月1日之前满足申领条件的用户都申领了一颗p004-胡杨，
     * 剩余的能量全部用来领取“p002-沙柳” 。
     * 统计在10月1日累计申领“p002-沙柳” 排名前10的用户信息；以及他比后一名多领了几颗沙柳。
     * 得到的统计结果如下表样式：
     */
    //过滤日期是2017年1月1日到2017年10月1日之间的
    userLowCarbonDF.where($"date_dt" >= "2017/1/1" and $"date_dt" <= "2017/10/1")//.show()
      //根据用户，日期分组，聚合每一天总的排放量
      .groupBy($"user_id")
      .agg(sum($"low_carbon") as "low_carbon")//.show()
      //新增一列，表示申领条件后的剩余能量
      .withColumn("other_carbon",when($"low_carbon" >= huYangCarbon,$"low_carbon" - huYangCarbon).otherwise($"low_carbon"))//.show()
      //新增一列，计算领取沙柳的棵树
      .withColumn("plant_count",floor($"other_carbon" / shaLiuCarbon))//.show()
      //新增一列，取出后一个沙柳的棵树
      .withColumn("after_plant_count",lead($"plant_count",1,0) over Window.orderBy($"plant_count".desc))
      .withColumn("less_count",$"plant_count" - $"after_plant_count")
      .limit(10)
      .select($"user_id",$"plant_count",$"less_count")
    //      .show()

    /**
     * 题目二：蚂蚁森林低碳用户排名分析
     *    查询user_low_carbon表中每日流水记录，条件为：
     * 用户在2017年，连续三天（或以上）的天数里，
     * 每天减少碳排放（low_carbon）都超过100g的用户低碳流水。
     * 需要查询返回满足以上条件的user_low_carbon表中的记录流水。
     * 例如用户u_002符合条件的记录如下，因为2017/1/2~2017/1/5连续四天的碳排放量之和都大于等于100g：
     */
    //根据用户和日期进行分组，得到每一天碳排放量
    userLowCarbonDF.groupBy($"user_id",$"date_dt")
      .agg(sum($"low_carbon") as "day_carbon")
      //过滤出大于100碳排放量的天
      .where($"day_carbon" > 100)
      //根据用户开窗，以日期升序排序
      .withColumn("rn",row_number() over Window.partitionBy($"user_id").orderBy($"date_dt"))
      //将日期减去编号，根据结果判断天数是否连续
      .withColumn("tmp_date",date_sub(regexp_replace($"date_dt","/","-"),$"rn"))
      //新增一列，计算用户连续的天数
      .withColumn("days",count(expr("1")) over Window.partitionBy($"user_id",$"tmp_date"))
      //过滤出连续天数是大于3的
      .where($"days" >= 3)
      .select($"user_id",$"date_dt")
      .join(userLowCarbonDF,List("user_id","date_dt"))
      .select($"user_id",$"date_dt",$"low_carbon")
      .show(1000)







  }
}